#!/bin/sh

set -e

# shellcheck disable=SC1091
. /usr/libexec/os-helpers-logging

BALENA_HOSTAPP_EXTENSIONS_FEATURE="io.balena.features.host-extension"
BALENA_HOSTAPP_EXTENSIONS_DEFAULT="/etc/hostapp-extensions.conf"
DOCKER_OVERLAY_ROOT="/var/lib/docker/overlay2/"

readonly script_name=$(basename "${0}")
usage() {
	cat <<EOF
Usage: ${script_name} [OPTIONS]
	-t Space separated list of repository tags for hostapp extension containers
	-r Reboot after update
EOF
	exit 0
}

while getopts 'hrt:' flag; do
	case "${flag}" in
	r) reboot=1 ;;
	t) target_hostext_images="${OPTARG}" ;;
	h) usage ;;
	*) error "Unexpected option ${flag}" ;;
	esac
done

shift $#

# Don't source before parsing args as balena-config-vars parses args too
# shellcheck disable=SC1091
. /usr/sbin/balena-config-vars

error_handler() {
	# shellcheck disable=SC2181
	[ $? -eq 0 ] && exit 0

	exit 1
}

trap error_handler EXIT

current_hostext_images=""
update_current_host_extensions() {
	if [ -z "${current_hostext_images}" ]; then
		current_hostext_images="${hostext}"
	else
		current_hostext_images="${current_hostext_images} ${hostext}"
	fi
}

umount_merged() {
	mergeddir=$(${DOCKER} inspect -f '{{.GraphDriver.Data.MergedDir}}' "${1}" 2> /dev/null) || mergeddir=""
	if [ -n "${mergeddir}" ]; then
		mergeddir=${mergeddir#${DOCKER_OVERLAY_ROOT}}
		overlay=$(findmnt --noheadings --canonicalize  -t overlay -l -o TARGET | grep "${mergeddir}" || true)
		if [ -n "${overlay}" ]; then
			umount -l "${overlay}" || true
		fi
	fi
}

# The target hostapp extensions is a space separated list of docker repository
# tags in the form repo/name:tag, e.g balena/default:v1 balena/extra:v2
#
# Select target hostapp extensions in order of preference:
# Command line argument
# hostappExtensions entry in config.json
# Image default from /etc/hostapp-extensions.conf
if [ -z "${target_hostext_images}" ]; then
	if [ -z "${HOSTEXT_IMAGES}" ]; then
		target_hostext_images=$(cat ${BALENA_HOSTAPP_EXTENSIONS_DEFAULT})
	else
		target_hostext_images="${HOSTEXT_IMAGES}"
	fi
fi
if [ -n "${target_hostext_images}" ]; then
	info "Target hostapp extensions:"
	for hostext in ${target_hostext_images}; do
		info " - ${hostext}"
	done
fi

# Detect containers engine
if which docker > /dev/null 2>&1; then
    DOCKER=docker
elif which rce > /dev/null 2>&1; then
    DOCKER=rce
elif which balena > /dev/null 2>&1; then
    DOCKER=balena
else
    error "No container engine detected."
    error_handler "no container engine detected"
fi

# Verify targets
for hostext in ${target_hostext_images}; do
	imageid=$(${DOCKER} inspect -f '{{.Id}}' "${hostext}" 2> /dev/null) || imageid=""
	if [ -n "${imageid}" ]; then
		continue
	fi

	if ! DOCKER_CLI_EXPERIMENTAL=enabled ${DOCKER} manifest inspect "${hostext}"  > /dev/null 2>&1; then
		error "${hostext} not found in repository - bailing out with no action taken."
		exit 1
	fi
done

info "Removing previous hostapp extensions"

for image in $(${DOCKER} images --all --quiet --filter label="${BALENA_HOSTAPP_EXTENSIONS_FEATURE}"); do
	imagetag=$(${DOCKER} inspect "${image}" --format='{{index .RepoTags 0}}' 2>/dev/null)
	cids=$(${DOCKER} ps --all --quiet --no-trunc --filter ancestor="${image}" | tr '\n' ' ')
	case "${target_hostext_images}" in
		*${imagetag}*)
			# If no container use the image do not preserve it
			if [ -n "${cids}" ]; then
				info " - Preserving $imagetag"
				continue
			fi
			;;
	esac
	for cid in ${cids}; do
		umount_merged "${cid}"
		${DOCKER} rm --force --volumes "${cid}" > /dev/null 2>&1 || true
	done
	${DOCKER} rmi --force "$image" > /dev/null 2>&1 || true
done

if [ -n "${target_hostext_images}" ]; then
	info "Installing new hostapp extensions:"
	for hostext in ${target_hostext_images}; do
		info " - ${hostext}"
		imageid=$(${DOCKER} inspect -f '{{.Id}}' "${hostext}" 2> /dev/null) || imageid=""
		if [ -n "${imageid}" ]; then
			info "   - already installed"
			update_current_host_extensions
			continue
		fi

		if ! ${DOCKER} pull "${hostext}"  > /dev/null 2>&1; then
			error "   - failed to install"
			break
		else
			if ! ${DOCKER} create --runtime="bare" --label "${BALENA_HOSTAPP_EXTENSIONS_FEATURE}" "${hostext}" none  > /dev/null 2>&1; then
				error "   - failed to spawn"
				break
			fi
			update_current_host_extensions
		fi
		info "   - done"
	done
fi

# Update config.json accordingly
tmpdir=$(mktemp -d)
tmpfile="${tmpdir}/config.tmp"
if [ -z "${current_hostext_images}" ]; then
	# No hostext images, remove entry
	jq . > "${tmpfile}" < "${CONFIG_PATH}"
	sed -i "/hostappExtensions/d" "${tmpfile}"
else
	# Store the current hostapp extensions into config.json
	jq -S '. |= .+ {"hostappExtensions"}' > "${tmpfile}" < "${CONFIG_PATH}"
	content=$(jq -S ". |= .+ {\"hostappExtensions\" : \"${current_hostext_images}\"}" < "${tmpfile}")
	echo "${content}" > "${tmpfile}"
fi
"${MV}" "${tmpfile}" "${CONFIG_PATH}"
rm -rf "${tmpdir}"

if [ -n "${current_hostext_images}" ]; then
	info "The following hostapp extensions will be available on the next boot:"
	for hostext in ${current_hostext_images}; do
		info " - ${hostext}"
	done
	if [ "$reboot" = 1 ]; then
		info "Rebooting now."
		reboot
	fi
else
	info "Nothing else to do"
fi

exit 0
